#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
认证处理器

该模块提供了各种HTTP认证处理器，支持基本认证、令牌认证、OAuth2认证等。

创建时间: 2023-07-15
作者: Python Tools Team
"""

import base64
import time
from typing import Dict, Any, Optional, Union, Callable, Tuple
import requests
from requests.auth import AuthBase

from .exception import HttpClientException


class BasicAuth(AuthBase):
    """
    基本认证处理器
    
    实现HTTP基本认证（Basic Authentication）。
    """
    
    def __init__(self, username: str, password: str):
        """
        初始化基本认证处理器
        
        Args:
            username: 用户名
            password: 密码
        """
        self.username = username
        self.password = password
    
    def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
        """
        为请求添加基本认证头
        
        Args:
            request: 准备好的请求对象
            
        Returns:
            requests.PreparedRequest: 添加了认证头的请求对象
        """
        auth_str = f"{self.username}:{self.password}"
        encoded = base64.b64encode(auth_str.encode('utf-8')).decode('utf-8')
        request.headers['Authorization'] = f"Basic {encoded}"
        return request


class TokenAuth(AuthBase):
    """
    令牌认证处理器
    
    实现基于令牌的认证，支持多种令牌类型。
    """
    
    def __init__(self, token: str, token_type: str = "Bearer"):
        """
        初始化令牌认证处理器
        
        Args:
            token: 认证令牌
            token_type: 令牌类型，默认为Bearer
        """
        self.token = token
        self.token_type = token_type
    
    def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
        """
        为请求添加令牌认证头
        
        Args:
            request: 准备好的请求对象
            
        Returns:
            requests.PreparedRequest: 添加了认证头的请求对象
        """
        request.headers['Authorization'] = f"{self.token_type} {self.token}"
        return request


class OAuth2Auth(AuthBase):
    """
    OAuth2认证处理器
    
    实现OAuth2认证，支持自动刷新令牌。
    """
    
    def __init__(self, 
                 access_token: str, 
                 token_type: str = "Bearer",
                 refresh_token: Optional[str] = None,
                 expires_at: Optional[float] = None,
                 refresh_callback: Optional[Callable[[str], Tuple[str, str, float]]] = None):
        """
        初始化OAuth2认证处理器
        
        Args:
            access_token: 访问令牌
            token_type: 令牌类型，默认为Bearer
            refresh_token: 刷新令牌，用于自动刷新访问令牌
            expires_at: 令牌过期时间戳
            refresh_callback: 刷新令牌的回调函数，接收刷新令牌作为参数，
                              返回新的(访问令牌, 刷新令牌, 过期时间戳)元组
        """
        self.access_token = access_token
        self.token_type = token_type
        self.refresh_token = refresh_token
        self.expires_at = expires_at
        self.refresh_callback = refresh_callback
    
    def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
        """
        为请求添加OAuth2认证头，如果令牌即将过期则尝试刷新
        
        Args:
            request: 准备好的请求对象
            
        Returns:
            requests.PreparedRequest: 添加了认证头的请求对象
            
        Raises:
            HttpClientException: 如果令牌刷新失败
        """
        # 如果令牌即将过期且有刷新机制，则刷新令牌
        if self.should_refresh():
            self.refresh()
        
        request.headers['Authorization'] = f"{self.token_type} {self.access_token}"
        return request
    
    def should_refresh(self) -> bool:
        """
        检查是否应该刷新令牌
        
        如果距离令牌过期时间不足60秒，且有刷新令牌和刷新回调，则返回True
        
        Returns:
            bool: 是否应该刷新令牌
        """
        return (self.expires_at is not None and 
                self.refresh_token is not None and 
                self.refresh_callback is not None and 
                time.time() > self.expires_at - 60)
    
    def refresh(self) -> None:
        """
        刷新访问令牌
        
        使用刷新令牌和刷新回调函数获取新的访问令牌
        
        Raises:
            HttpClientException: 如果令牌刷新失败
        """
        if not self.refresh_token or not self.refresh_callback:
            raise HttpClientException(
                message="无法刷新令牌：缺少刷新令牌或刷新回调函数",
                code="E100400"
            )
        
        try:
            self.access_token, self.refresh_token, self.expires_at = self.refresh_callback(self.refresh_token)
        except Exception as e:
            raise HttpClientException(
                message="令牌刷新失败",
                code="E100401",
                cause=e
            )